import torch
from torch.nn import CrossEntropyLoss


def get_f1(labels_list, preds, out_label_ids):
    gt_dict = {i: 0 for i in range(len(labels_list))}
    pred_dict = {i: 0 for i in range(len(labels_list))}
    correct_dict = {i: 0 for i in range(len(labels_list))}
    for i in range(len(preds)):
        current_label = torch.tensor(out_label_ids[i])

        re_label = torch.tensor(preds[i])
        idx = torch.where(current_label != CrossEntropyLoss().ignore_index)[0]
        current_label_no_pad = current_label[idx]
        re_label_no_pad = re_label[idx]
        pre_cur, pre_re = 0, 0
        cur_spans, re_spans = [], []
        for j in range(len(current_label_no_pad)):
            if current_label_no_pad[j] != 0:
                if pre_cur == 0:
                    start = j
                elif pre_cur != 0 and pre_cur != current_label_no_pad[j]:
                    end = j
                    gt_dict[pre_cur.item()] += 1
                    cur_spans.append((start, end, pre_cur))

                pre_cur = current_label_no_pad[j]
            else:
                if pre_cur != 0:
                    end = j
                    gt_dict[pre_cur.item()] += 1
                    cur_spans.append((start, end, pre_cur))
                pre_cur = 0

            if re_label_no_pad[j] != 0:
                if pre_re == 0:
                    start_ = j
                elif pre_re != 0 and pre_re != re_label_no_pad[j]:
                    end_ = j

                    pred_dict[pre_re.item()] += 1
                    re_spans.append((start_, end_, pre_re))
                pre_re = re_label_no_pad[j]
            else:
                if pre_re != 0:
                    end_ = j
                    pred_dict[pre_re.item()] += 1
                    re_spans.append((start_, end_, pre_re))
                pre_re = 0
        if len(cur_spans) != 0:
            for cur in cur_spans:
                for re in re_spans:
                    if cur[0] == re[0] and cur[1] == re[1] and cur[2] == re[2]:
                        correct_dict[cur[2].item()] += 1
                        break

    precision_dict = {i: 0 for i in range(len(labels_list))}
    recall_dict = {i: 0 for i in range(len(labels_list))}
    type_f1_dict = {i: 0 for i in range(len(labels_list))}
    for i in range(len(labels_list)):
        if pred_dict[i] != 0:
            precision_dict[i] = round(correct_dict[i] / pred_dict[i], 2)
        if gt_dict[i] != 0:
            recall_dict[i] = round(correct_dict[i] / gt_dict[i], 2)
        if precision_dict[i] != 0 and recall_dict[i] != 0:
            type_f1_dict[i] = round(2 * precision_dict[i] * recall_dict[i] / (precision_dict[i] + recall_dict[i]), 2)

    marco_f1 = sum(type_f1_dict.values()) / (len(labels_list) - 1)


    print("gt_dict: ", gt_dict)
    print("pred_dict: ", pred_dict)
    print("correct_dict: ", correct_dict)
    print("precision_dict: ", precision_dict)
    print("recall_dict: ", recall_dict)
    total_precision = sum(correct_dict.values()) / sum(pred_dict.values())
    total_recall = sum(correct_dict.values()) / sum(gt_dict.values())
    total_f1 = 2 * total_precision * total_recall / (total_precision + total_recall)
    print("type_f1_dict: ", type_f1_dict)
    print("marco_f1: ", marco_f1)
    print("total_f1: ", total_f1)




